"""
自动填充模块

这个模块实现了自动填充功能，主要用于处理序列长度需要是特定倍数的情况。
主要特点：
1. 自动检测和填充序列到指定倍数长度
2. 处理输入掩码的相应填充
3. 支持自定义填充值
4. 保持输出序列的原始长度
"""

import math
import torch
from torch import nn
from routing_transformer.routing_transformer import RoutingTransformer
import torch.nn.functional as F

def find_module(nn_module, type):
    """
    在神经网络模块中查找特定类型的子模块
    
    参数:
        nn_module: PyTorch模块
        type: 要查找的模块类型
    
    返回:
        找到的第一个匹配模块，如果没找到则返回None
    """
    for module in nn_module.modules():
        if isinstance(module, type):
            return module
    return None

def pad_to_multiple(tensor, multiple, dim=-1, value=0):
    """
    将张量在指定维度填充到multiple的整数倍
    
    参数:
        tensor: 输入张量
        multiple: 目标倍数
        dim: 需要填充的维度（默认为最后一维）
        value: 填充值（默认为0）
    
    返回:
        填充后的张量，长度为multiple的整数倍
    """
    seqlen = tensor.shape[dim]
    m = seqlen / multiple
    if m.is_integer():
        return tensor

    pre_pad_offset = (0,) * (-1 - dim) * 2
    padding = math.ceil(m) * multiple - seqlen
    padded_tensor = F.pad(tensor, (*pre_pad_offset, *(0, padding)), value=value)
    return padded_tensor

class Autopadder(nn.Module):
    """
    自动填充包装器
    
    为模型提供自动填充功能的包装器，确保输入序列长度满足模型要求。
    主要功能：
    1. 自动检测所需的填充倍数
    2. 处理输入序列和掩码的填充
    3. 在输出时恢复原始序列长度
    
    参数:
        net: 要包装的网络模块
    """
    def __init__(self, net):
        super().__init__()
        # 查找Transformer模块并获取填充倍数
        transformer = find_module(net, RoutingTransformer)
        self.net = net
        self.pad_multiple = transformer.pad_to_multiple

    def forward(self, x, **kwargs):
        """
        前向传播
        
        参数:
            x: 输入张量
            **kwargs: 其他参数，可能包含input_mask
        
        返回:
            (output, loss): 模型输出和损失
            - output的长度与输入序列相同
            - 自动处理了填充和掩码
        """
        # 如果不需要填充，直接返回原始结果
        if self.pad_multiple <= 0:
            return self.net(x, **kwargs)

        b, t, device = *x.shape, x.device

        # 获取或创建输入掩码
        input_mask = kwargs.get('input_mask')
        if input_mask is None:
            input_mask = torch.full((b, t), True, device=device, dtype=torch.bool)

        # 对输入和掩码进行填充
        x = pad_to_multiple(x, self.pad_multiple, dim=1)
        new_mask = pad_to_multiple(input_mask, self.pad_multiple, dim=1, value=False)
        kwargs.update(input_mask=new_mask)

        # 运行模型并截取回原始长度
        out, loss = self.net(x, **kwargs)
        return out[:, 0:t], loss
